#include "znet.h"

#include <stdint.h>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <sys/epoll.h>


/*
 * sock: session socket fd
 * epfd: epoll fd
 */
static zsess_t *sess_new(znet_t *srv, int sock)
{
    /* create client object */
    zsess_t *sess = (zsess_t *)malloc(sizeof(zsess_t));
    if (sess == NULL) {
        printf("malloc failed!\n");
        return NULL;
    }
    memset(sess, 0, sizeof(zsess_t));

    sess->sock = sock;
    list_insert(&srv->head, &sess->node);

    /* add client to epoll */
    struct epoll_event ev = {
        .events = EPOLLIN | EPOLLRDHUP | EPOLLERR,
        .data.ptr = sess,
    };

    if (epoll_ctl(srv->epfd, EPOLL_CTL_ADD, sock, &ev) < 0) {
        printf("epoll add: failed!\n");
        return NULL;
    }

    return sess;
}

static void sess_del(znet_t *srv, zsess_t *sess)
{
    epoll_ctl(srv->epfd, EPOLL_CTL_DEL, sess->sock, NULL);
    list_remove(&sess->node);
    free(sess);
    sess = NULL;
}


static int client_init(znet_t *cli)
{
    int ret = 0;

    /* get socket */
    cli->sock = socket(AF_INET, SOCK_STREAM, 0);
    if (cli->sock < 0) {
        printf("socket: failed!\n");
        return -1;
    }

    struct sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_addr.s_addr = inet_addr(cli->addr),
        .sin_port = htons(cli->port),
    };

    socklen_t len = sizeof(struct sockaddr_in);

    /* connect socket */
    ret = connect(cli->sock, (const struct sockaddr *)&addr, len);
    if (ret < 0) {
        printf("connect: failed!\n");
        return -1;
    }

    return 0;
}

static int server_init(znet_t *srv)
{
    int ret = 0;

    struct sockaddr_in addr = {
        .sin_family      = AF_INET,
        .sin_addr.s_addr = htonl(INADDR_ANY),
        .sin_port        = htons(srv->port),
    };

    /* get socket */
    srv->sock = socket(AF_INET, SOCK_STREAM, 0);
    if (srv->sock < 0) {
        printf("socket: failed!\n");
        return -1;
    }

    /* reuse socket addr */
    int opt = SO_REUSEADDR;
    setsockopt(srv->sock, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    /* bind socket */
    ret = bind(srv->sock, (struct sockaddr *)&addr, sizeof(addr));
    if (ret < 0) {
        printf("bind: failed!\n");
        close(srv->sock);
        return ret;
    }

    /* if port is 0, get real port */
    if (srv->port == 0) {
        socklen_t addrsize = sizeof(addr);
        getsockname(srv->sock, (struct sockaddr *)&addr, &addrsize);
        srv->port = ntohs(addr.sin_port);
    }

    /* listen socket */
    ret = listen(srv->sock, ZNET_TCP_BACKLOG);
    if (ret < 0) {
        printf("listen: failed!\n");
        close(srv->sock);
        return ret;
    }

    return ret;
}

static int server_epoll(znet_t *srv)
{
    int epfd;
    struct epoll_event events[ZNET_TCP_SESSION_MAX];

    /* add server to epoll */
    epfd = epoll_create(1);
    srv->epfd = epfd;

    struct epoll_event ev = {
        .events   = EPOLLIN,
        .data.ptr = srv,
    };

    if (epoll_ctl(epfd, EPOLL_CTL_ADD, srv->sock, &ev) < 0) {
        printf("epoll: add error!\n");
        return -1;
    }

    /* temp pointer for client */
    zsess_t *sess = NULL;
    list_init(&srv->head);

    for (;;) {
        /* epoll wait */
        int nepoll = epoll_wait(epfd, events, ZNET_TCP_SESSION_MAX, -1);
        if (nepoll < 0) {
            printf("epoll: wait error!\n");
            return -1;
        }

        for (int i = 0; i < nepoll; i++) {
            if (events[i].data.ptr == srv) { /* is server sock wake up, a new client connected. */
                /* get client sock */
                struct sockaddr_in addr;
                socklen_t len = sizeof(addr);

                int sock = accept(srv->sock, (struct sockaddr *)&addr, &len);
                if (sock < 0) {
                    printf("accept: failed!\n");
                    return -1;
                }

                /* create a new session, will add it to epoll and server list. */
                sess = sess_new(srv, sock);
                inet_ntop(AF_INET, &addr.sin_addr.s_addr, sess->addr, sizeof(sess->addr));
                sess->port = ntohs(addr.sin_port);
                if (srv->on_open(sess) < 0) {
                    close(sess->sock);
                }
            } else { /* is client socket wake up, client data is comming. */
                sess = events[i].data.ptr;
                int event = events[i].events;
                if (event & EPOLLIN) {
                    srv->on_recv(sess);
                }

                if (event & EPOLLRDHUP || event & EPOLLERR) {
                    printf("err: %d, rdhup: %d\n", !!(event & EPOLLERR), !!(event & EPOLLRDHUP));
                    srv->on_exit(sess);
                    close(sess->sock);
                    /* delete sess, from epoll and server list. */
                    sess_del(srv, sess);
                    events[i].data.ptr = NULL;
                }
            }
        }
    } // end inf loop
}


int znet_launch(znet_t *znet, int role)
{
    int ret = 0;
    if (role == ZNET_SERVER) {
        if (znet->on_open == NULL || znet->on_recv == NULL || znet->on_exit == NULL) {
            printf("server: undefine callback!\n");
            return -1;
        }

        ret = server_init(znet);
        if (ret < 0) {
            return ret;
        }

        if (znet->on_launch != NULL) {
            znet->on_launch(znet->port);
        }
        return server_epoll(znet);
    } else if (role == ZNET_CLIENT) {
        return client_init(znet);
    } 
    return -1;
}

int znet_shutdown(int sock)
{
    return shutdown(sock, SHUT_RDWR);
}

int znet_send(int sock, void *buf, int len)
{
    return send(sock, buf, len, 0);
}

int znet_recv(int sock, void *buf, int len)
{
    return recv(sock, buf, len, 0);
}

int znet_send_all(int sock, void *buf, int len)
{
    int ret = 0;
    int off = 0;

redo:
    ret = send(sock, buf + off, len - off, 0);
    if (ret <= 0) {
        return ret;
    }
    off += ret;
    if (off < len) {
        goto redo;
    }

    return off;
}

int znet_recv_all(int sock, void *buf, int len)
{
    int ret = 0;
    int off = 0;

redo:
    ret = recv(sock, buf + off, len - off, 0);
    if (ret <= 0) {
        return ret;
    }
    off += ret;
    if (off < len) {
        goto redo;
    }

    return off;
}
